import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from scipy import ndimage
import h5py

class MultiImgs:
    def __init__(self,image_set=None,**kwargs):
        self.image_set=image_set
        self.image_sorted=self.sort()
        self.layers=self.image_sorted.keys()
        self.nlayers=len(self.layers)
    def list_layers(self):
        n=0
        layers=[]
        for x in self.image_set:
            if x.layer not in layers:
                layers.append(x.layer)
                n+=1
        return layers
    def sort(self):
        image_sorted={}
        for x in self.image_set:
            if x.layer in image_sorted.keys():
                image_sorted[x.layer].append(x)
            else:
                image_sorted[x.layer]=[x]
        return image_sorted
    def add_image(self,new_image):
        self.image_set.append(new_image)
        if new_image.layer in self.image_sorted.keys():
            self.image_sorted[new_image.layer].append(new_image)
        else:
            self.image_sorted[new_image.layer]=[new_image]
    def show(self,**kwargs):
        show_layers=kwargs.get('show_layers',self.layers)
        try:
            show_layers=sorted(show_layers)
        except TypeError:
            show_layers=[show_layers]
        extent=kwargs.get('extent',self.image_sorted[self.layers[0]][0].axis)
        vmin=kwargs.get('vmin')
        vmax=kwargs.get('vmax')
        for layer in show_layers:
            for image in self.image_sorted[layer]:
                if layer==show_layers[0] and image==self.image_sorted[layer][0]:
                    if not vmin:
                        kwargs['vmin']=np.nanmin(image.image)
                    if not vmax:
                        kwargs['vmax']=np.nanmax(image.image)
                    image.show(hold=True,**kwargs)
                else:
                    image.show(hold=True,vmin=vmin,vmax=vmax,scalebar='off',colorbar='off')
        plt.xlim(extent[:2])
        plt.ylim(extent[-2:])
        plt.show()
#    def out(self,**kwargs):
#        extent=kwargs.get('extent',self.image_sorted[self.layers[0]][0].axis)
#        layers=kwargs.get('layers',self.layers)
        
    
class MultiImg:
    def __init__(self,image=None,layer=None,**kwargs):
        self.layer=layer
        if type(image) is type(mapsimg()):
            for attr in dir(image):
                exec('self.'+attr+'=image.'+attr)
        else:
            self.image=image
            self.axis=kwargs.get('axis')
            self.name=kwargs.get('name')
            self.unit=kwargs.get('unit')
            self.xcoord=kwargs.get('xcoord')
            self.ycoord=kwargs.get('ycoord')

class mapsimg:
    def __init__(self,image=None,name=None,unit=None,xaxis=None,yaxis=None,xcoord='mm',ycoord='mm'):
        self.image=image
        self.name=name
        self.unit=unit
        self.xaxis=xaxis
        self.yaxis=yaxis
        self.xcoord=xcoord
        self.ycoord=ycoord
        try:
            self.axis=[xaxis[0],xaxis[-1],yaxis[0],yaxis[-1]]
        except TypeError:
            self.axis=None
    def loadh5(self,filename,element=None):
        self.image,self.name,self.unit,self.xaxis,self.yaxis,self.xcoord,self.ycoord=readh5(filename,element)
        self.axis=[self.xaxis[0],self.xaxis[-1],self.yaxis[0],self.yaxis[-1]]
    def show(self,**kwargs):
        element=kwargs.get('element',None)
        axis=kwargs.get('axis','on')
        scale_bar=kwargs.get('scalebar','on')
        colorbar=kwargs.get('colorbar','on')
        hold=kwargs.get('hold',False)
        vmin=kwargs.get('vmin')
        vmax=kwargs.get('vmax')
        extent=kwargs.get('extent',self.axis)
        if element!=None:
            n=self.search(element)
            image=self.image[n]
        else:
            image=self.image
        plt.imshow(image,extent=self.axis,aspect='equal',vmin=vmin,vmax=vmax)
        if axis=='off':
            plt.axis('off')
        if colorbar=='on':
            plt.colorbar()
        if scale_bar=='on':
            scalebar(axis=extent,unit=self.xcoord,**kwargs)
        plt.xlabel(self.xcoord)
        plt.ylabel(self.ycoord)
        if hold:
            plt.hold(True)
        else:
            plt.xlim(extent[:2])
            plt.ylim(extent[-2:])
            plt.show()
    def search(self,element):
        n=np.where(self.name==element)[0][0]
        return n



def scalebar(**kwargs):
    """Add scalebar (:class:`~matplotlib.patches.Patch`) on current axis 
    (:class:`~matplotlib.axes.Axes`).
    
    Keyword arguments:
        *axis*: scalars (left, right, bottom, top)
            Frame limits on current axis. If given, the function will take the 
            region delineated by axis as the region to display. If not given, 
            the default value is the current axis given by plt.gca()
        
        *loc*: ['lowerleft' | 'lowerright' | 'upperleft' | 'upperright']
            Set the location of the scalebar on the frame. Default='lowerleft'.
        
        *color*: any matplotlib color
            Set the color of scalebar and text on scalebar. Default='black'
        
        *textcolor*: any matplotlib color
            Set the color of text on scalebar. Default=color
        
        *barcolor*: any matplotlib color
            Set the color of scalebar. Default=color
        
        *unit*: string
            Unit used in current axis. Default='mm'
        
        *barlen*: scaler
            Length of the scalebar in unit designated in 'unit'. Default=1
            
        *direction*: ['horizontal' | 'vertical']
            Direction of the scalebar. Default='horizontal' 
            
        *text_direction*: ['horizontal' | 'vertical']
            Direction of the text on scalebar. Default='horizontal' 
            
        *texton*: [True | False]
            Show text on scalebar or not. Default=True
        
    **Example**
        import numpy as np
        import matplotlib.pyplot as plt
        x=np.arange(100)
        y=np.arange(100)
        x,y=np.meshgrid(x,y)
        img=x*y
        axis=[-5.2,-15.1,-1.3,-6.4]
        plt.imshow(img,origin='lower',extent=axis)
        scalebar(color='white',unit='cm',barlen=2)
        plt.show()
        """
    
    plt.hold(True)
    axis=kwargs.get('axis',plt.axis())
    loc=kwargs.get('loc','lowerleft')
    color=kwargs.get('color','black')
    textcolor=kwargs.get('textcolor',color)
    barcolor=kwargs.get('barcolor',color)
    unit=kwargs.get('unit','mm')
    direction=kwargs.get('direction','horizontal')
    text_direction=kwargs.get('text_direction',direction)
    barlen=kwargs.get('barlen',1)
    texton=kwargs.get('texton',True)


    #scalebar location
    if loc=='lowerleft':
        xloc=axis[0]+(axis[1]-axis[0])*0.1
        yloc=axis[2]+(axis[3]-axis[2])*0.1
        islower=1
        isleft=1
    elif loc=='upperleft':
        xloc=axis[0]+(axis[1]-axis[0])*0.1
        yloc=axis[3]-(axis[3]-axis[2])*0.1
        islower=-1
        isleft=1
    elif loc=='lowerright':
        xloc=axis[1]-(axis[1]-axis[0])*0.1
        yloc=axis[2]+(axis[3]-axis[2])*0.1
        islower=1
        isleft=-1
    elif loc=='upperright':
        xloc=axis[1]-(axis[1]-axis[0])*0.1
        yloc=axis[3]-(axis[3]-axis[2])*0.1
        islower=-1
        isleft=-1
    elif type(loc) is not str:
        xloc=loc[0]
        yloc=loc[1]
        islower=1
        isleft=1
    #scalebar length
    x_axis_direction=np.sign(axis[1]-axis[0])
    y_axis_direction=np.sign(axis[3]-axis[2])
    if direction=='vertical':
        ylen=barlen*y_axis_direction
        xlen=(axis[1]-axis[0])*0.02
        if texton:
            x_text=xloc+isleft*xlen*2.2
            y_text=yloc+islower*ylen/2. 
            plt.text(x_text,y_text,str(barlen)+' '+unit,rotation=text_direction,color=textcolor,verticalalignment='center',withdash=True,horizontalalignment='center')     
    else:
        xlen=barlen*x_axis_direction
        ylen=(axis[3]-axis[2])*0.02
        if texton:
            x_text=xloc+isleft*xlen/2.
            y_text=yloc+islower*ylen*2.2 
            plt.text(x_text,y_text,str(barlen)+' '+unit,rotation=text_direction,color=textcolor,verticalalignment='center',withdash=True,horizontalalignment='center')   
    current_axis=plt.gca()
    current_axis.add_patch(Rectangle([xloc,yloc],isleft*xlen,islower*ylen,color=barcolor))
    

def loadh5(filename,element=None):
    f=h5py.File(filename,'r')
    name=np.array(f['exchange']['images_names'])
    xaxis=list(f['exchange']['x_axis'])[::-1]
    yaxis=list(f['exchange']['y_axis'])[::-1]
    n=np.where(name=='x_coord')[0][0]
    xcoord=f['exchange']['images_units'][n]
    n=np.where(name=='y_coord')[0][0]
    ycoord=f['exchange']['images_units'][n]
    if element!=None:
        n=np.where(name==element)[0][0]
        image=f['exchange']['images'][n]
        name=name[n]
        unit=f['exchange']['images_units'][n]
    else:
        image=list(f['exchange']['images'])
        unit=list(f['exchange']['images_units'])
    f.close()
    return mapsimg(image,name,unit,xaxis,yaxis,xcoord,ycoord)

def readh5(filename,element=None):
    f=h5py.File(filename,'r')
    name=np.array(f['exchange']['images_names'])
    xaxis=list(f['exchange']['x_axis'])[::-1]
    yaxis=list(f['exchange']['y_axis'])[::-1]
    n=np.where(name=='x_coord')[0][0]
    xcoord=f['exchange']['images_units'][n]
    n=np.where(name=='y_coord')[0][0]
    ycoord=f['exchange']['images_units'][n]
    if element!=None:
        n=np.where(name==element)[0][0]
        image=f['exchange']['images'][n]
        name=name[n]
        unit=f['exchange']['images_units'][n]
    else:
        image=list(f['exchange']['images'])
        unit=list(f['exchange']['images_units'])
    f.close()
    return image,name,unit,xaxis,yaxis,xcoord,ycoord

def fftfreq2d(m,n=None,d=1,zero='center',axis=0):
    x=np.fft.fftfreq(m,d)
    if n==None: n=m
    y=np.fft.fftfreq(n,d)
    X,Y=np.meshgrid(x,y)
    if axis==0: result=X
    else: result=Y
    if zero=='center':
        result=np.fft.fftshift(result,1-axis)
    return result

def shiftimg(im_fft,shifts,zero='front'):
    """fft of image -> shifted fft of image by x=shifts[0], y=shifts[1]"""
    im_fft_shift=im_fft[:]
    for n in range(len(shifts)):
        im_fft_shift=im_fft_shift*np.exp(1j*fftfreq2d(np.shape(im_fft)[1],np.shape(im_fft)[0],axis=n,zero=zero)*2*np.pi*(-shifts[n]))
    return im_fft_shift

def zeropadding(im,s):
    """Pad ndarray im with zeros arround the original ndarray.
    Number of zeros (on each side) padded for different dimensions are declared in s.
    If the length of s is shorter than the dimensions of im, s will be padded with
    its last value in the end."""
    dim=len(np.shape(im))
    if type(s)==int:
        s=[s]
    if len(s)<dim:
        s0=s[-1]
        for n in range(len(s),dim):
            s=np.append(s,s0)
    im_padded=im[:]
    for n in range(dim):
        shape=list(np.shape(im_padded))
        shape[n]=s[n]
        append=np.zeros(shape)
        im_padded=np.append(append,im_padded,axis=n)
        im_padded=np.append(im_padded,append,axis=n)
    return im_padded

